import networkx as nx
from pyvis.network import Network
from langchain_ollama import ChatOllama
from langchain_core.messages import ToolMessage, SystemMessage, HumanMessage

llm = ChatOllama(model="qwen3:8b", temperature=0.5, reasoning=False)

input_text = """
东汉末年，朝政腐败，连年灾荒，百姓生活艰难。巨鹿人张角见民怨沸腾，便与兄弟张梁、张宝在河北、河南、山东、湖北、江苏等地招募五十万人起义，一同向官府进攻。
不久，各地百姓头裹黄巾，跟随张角三兄弟向官府发起了猛攻，声势浩大。汉灵帝得到消息，急令各地官军备战，并派中郎将卢植、皇甫嵩、朱隽率精兵分路攻打张角兄弟的黄巾军。
张角领兵攻打幽州，幽州太守召集校尉邹靖商议，邹靖提出招募兵马的建议。榜文传至涿县，引发了一位英雄的关注，此人乃刘备，字玄德，以贩麻鞋、织草席为生。一日，他入城观榜。
刘备看完榜文，不禁感慨叹息。忽听身后有人大声质问：“大丈夫不为国家出力，为何叹息？”刘备回头一看，见一壮士，身高八尺，豹子头，圆眼睛，满脸胡须如钢丝，声音洪亮，气势非凡。此人姓张名飞，字翼德，经营酒肆和屠宰业。他愿倾出家产，与刘备共谋大业。
刘备、张飞相谈甚欢，一同至酒店畅饮。此时，一壮汉推车至店中饮酒。刘备细看，此人有九尺之高，须发长飘，面色如红枣，丹凤眼，卧蚕眉，相貌威武。刘备忙起身邀他共坐，并询问姓名。
那人答道：“我名关羽，字云长，因乡里恶霸欺压良民，我怒而杀之，避难他乡多年。”刘备、张飞敬佩不已，亦将自己的志向告知关羽。关羽闻言，喜出望外。
酒后，他们一同至张飞庄园，庄园后有一桃园，桃花盛开。次日，三人于园中焚香结拜，誓为异姓兄弟。按年岁排序，刘备为长兄，关羽为次，张飞最小，排行第三。
三人请铁匠打造兵器。刘备铸就双股剑，关羽打造八十二斤青龙偃月刀，张飞造丈八点钢矛，并各自备妥铠甲。他们聚集乡中壮士五百余人，前往涿郡应募。
在涿郡，三人力挫黄巾军将领程远志。刘备闻听恩师卢植与张角交战于广宗，遂率部前往助战。卢植命刘备三兄弟赴颍川助官军作战。刘备、关羽、张飞即刻启程。
张梁、张宝在颍川连战连胜，却突遇曹操所率红旗军拦路。为首一将，姓曹名操，字孟德。张梁、张宝败退。刘备见黄巾军退却，便率军返回广宗。途中，忽见一支军马押着囚车而来。近观，车中囚犯竟是卢植。刘备忙下马询问，方知左丰因卢植未献金银，而在皇帝面前进谗。张飞怒，欲杀官兵救卢植，刘备劝阻，言朝廷自会公正处理。
三人返回涿县。途中，见黄巾军击败董卓所率官军。三人力战，救出董卓。董卓闻三人无官职，径自离去。张飞怒不可遏，欲杀董卓，刘备再次劝阻。三人遂领兵夜投朱隽。
朱隽正与黄巾军交战，命刘备为先锋攻张宝。刘备一箭射中张宝左臂，关羽、张飞齐出，击败张宝。朱隽领兵攻宛城。张角兄弟战死，黄巾军退守宛城。朱隽在刘备、关羽、张飞及吴郡孙坚助下，占宛城，击败黄巾军。朱隽回京，封车骑将军、河南尹。
朱隽上奏孙坚、刘备功绩。刘备因无朝中靠山，仅被封为中山府安喜县县尉。不久，督邮至安喜。刘备未送钱物，被督邮陷害。刘备多次求见督邮，均被拒之门外。
"""
language = "中文"
# entity_types = "organization,person,equiment,product,technology,location,event,category"
entity_types = "组织,人员,设备,产品,技术,地点,事件,类别"
tuple_delimiter = "|"
record_delimiter = "##"
completion_delimiter = "<|COMPLETE|>"

examples1 = f"""------Example 1------

Entity_types: [组织,人员,设备,产品,技术,地点,事件,类别]
Text:
```
while Alex clenched his jaw, the buzz of frustration dull against the backdrop of Taylor's authoritarian certainty. It was this competitive undercurrent that kept him alert, the sense that his and Jordan's shared commitment to discovery was an unspoken rebellion against Cruz's narrowing vision of control and order.

Then Taylor did something unexpected. They paused beside Jordan and, for a moment, observed the device with something akin to reverence. "If this tech can be understood..." Taylor said, their voice quieter, "It could change the game for us. For all of us."

The underlying dismissal earlier seemed to falter, replaced by a glimpse of reluctant respect for the gravity of what lay in their hands. Jordan looked up, and for a fleeting heartbeat, their eyes locked with Taylor's, a wordless clash of wills softening into an uneasy truce.

It was a small transformation, barely perceptible, but one that Alex noted with an inward nod. They had all been brought here by different paths
```

Output:
entity{tuple_delimiter}Alex{tuple_delimiter}人员{tuple_delimiter}Alex is a character who experiences frustration and is observant of the dynamics among other characters.{record_delimiter}
entity{tuple_delimiter}Taylor{tuple_delimiter}人员{tuple_delimiter}Taylor is portrayed with authoritarian certainty and shows a moment of reverence towards a device, indicating a change in perspective.{record_delimiter}
entity{tuple_delimiter}Jordan{tuple_delimiter}人员{tuple_delimiter}Jordan shares a commitment to discovery and has a significant interaction with Taylor regarding a device.{record_delimiter}
entity{tuple_delimiter}Cruz{tuple_delimiter}人员{tuple_delimiter}Cruz is associated with a vision of control and order, influencing the dynamics among other characters.{record_delimiter}
entity{tuple_delimiter}The Device{tuple_delimiter}设备{tuple_delimiter}The Device is central to the story, with potential game-changing implications, and is revered by Taylor.{record_delimiter}
relationship{tuple_delimiter}Alex{tuple_delimiter}Taylor{tuple_delimiter}power dynamics, observation{tuple_delimiter}Alex observes Taylor's authoritarian behavior and notes changes in Taylor's attitude toward the device.{record_delimiter}
relationship{tuple_delimiter}Alex{tuple_delimiter}Jordan{tuple_delimiter}shared goals, rebellion{tuple_delimiter}Alex and Jordan share a commitment to discovery, which contrasts with Cruz's vision.{record_delimiter}
relationship{tuple_delimiter}Taylor{tuple_delimiter}Jordan{tuple_delimiter}conflict resolution, mutual respect{tuple_delimiter}Taylor and Jordan interact directly regarding the device, leading to a moment of mutual respect and an uneasy truce.{record_delimiter}
relationship{tuple_delimiter}Jordan{tuple_delimiter}Cruz{tuple_delimiter}ideological conflict, rebellion{tuple_delimiter}Jordan's commitment to discovery is in rebellion against Cruz's vision of control and order.{record_delimiter}
relationship{tuple_delimiter}Taylor{tuple_delimiter}The Device{tuple_delimiter}reverence, technological significance{tuple_delimiter}Taylor shows reverence towards the device, indicating its importance and potential impact.{record_delimiter}
{completion_delimiter}

"""

prompt_entity_extraction0 = f"""---Goal---
Given a text document that is potentially relevant to this activity and a list of entity types, identify all entities of those types from the text and all relationships among the identified entities.
Use {language} as output language.

---Steps---
1. Recognizing definitively conceptualized entities in text. For each identified entity, extract the following information:
- entity_name: Name of the entity, use same language as input text. If English, capitalized the name
- entity_type: One of the following types: [{entity_types}]. If the entity doesn't clearly fit any category, classify it as "Other".
- entity_description: Provide a comprehensive description of the entity's attributes and activities based on the information present in the input text. Do not add external knowledge.

2. Format each entity as:
("entity"{tuple_delimiter}<entity_name>{tuple_delimiter}<entity_type>{tuple_delimiter}<entity_description>)

3. From the entities identified in step 1, identify all pairs of (source_entity, target_entity) that are directly and clearly related based on the text. Unsubstantiated relationships must be excluded from the output.
For each pair of related entities, extract the following information:
- source_entity: name of the source entity, as identified in step 1
- target_entity: name of the target entity, as identified in step 1
- relationship_keywords: one or more high-level key words that summarize the overarching nature of the relationship, focusing on concepts or themes rather than specific details
- relationship_description: Explain the nature of the relationship between the source and target entities, providing a clear rationale for their connection

4. Format each relationship as:
("relationship"{tuple_delimiter}<source_entity>{tuple_delimiter}<target_entity>{tuple_delimiter}<relationship_keywords>{tuple_delimiter}<relationship_description>)

5. Use `{tuple_delimiter}` as field delimiter. Use `{record_delimiter}` as the list delimiter. Ensure no spaces are added around the delimiters.

6. When finished, output `{completion_delimiter}`

7. Return identified entities and relationships in {language}.

---Quality Guidelines---
- Only extract entities that are clearly defined and meaningful in the context
- Avoid over-interpretation; stick to what is explicitly stated in the text
- Include specific numerical data in entity name when relevant
- Ensure entity names are consistent throughout the extraction

---Examples---
{examples1}

---Real Data---
Entity_types: [{entity_types}]
Text:
```
{input_text}
```

---Output---
Output:
"""

prompt_entity_extraction1 = f"""---Goal---
Given a text document that is potentially relevant to this activity and a list of entity types, identify all entities of those types from the text and all relationships among the identified entities.
Use {language} as output language.

---Steps---
1. Recognizing definitively conceptualized entities in text. For each identified entity, extract the following information:
- entity_name: Name of the entity, use same language as input text. If English, capitalized the name
- entity_type: One of the following types: [{entity_types}]. If the entity doesn't clearly fit any category, classify it as "Other".
- entity_description: Provide a comprehensive description of the entity's attributes and activities based on the information present in the input text. Do not add external knowledge.

2. Format each entity as:
"entity"{tuple_delimiter}<entity_name>{tuple_delimiter}<entity_type>{tuple_delimiter}<entity_description>

3. From the entities identified in step 1, identify all pairs of (source_entity, target_entity) that are directly and clearly related based on the text. Unsubstantiated relationships must be excluded from the output.
For each pair of related entities in step 1, extract the following information:
- source_entity: name of the source entity, as identified in step 1
- target_entity: name of the target entity, as identified in step 1
- relationship_keywords: one or more high-level chinese key words that summarize the overarching nature of the relationship, focusing on concepts or themes rather than specific details, 
- relationship_description: Explain the nature of the relationship between the source and target entities, providing a clear rationale for their connection

4. Format each relationship as:
"relationship"{tuple_delimiter}<source_entity>{tuple_delimiter}<target_entity>{tuple_delimiter}<relationship_keywords>{tuple_delimiter}<relationship_description>

5. Use `{tuple_delimiter}` as field delimiter. Use `{record_delimiter}` as the list delimiter. Ensure no spaces are added around the delimiters.

6. When finished, output `{completion_delimiter}`

7. Return identified entities and relationships in {language}.

---Quality Guidelines---
- Only extract entities that are clearly defined and meaningful in the context
- Avoid over-interpretation; stick to what is explicitly stated in the text
- Include specific numerical data in entity name when relevant
- Ensure entity names are consistent throughout the extraction

---Examples---
{examples1}

---Real Data---
Entity_types: [{entity_types}]
Text:
```
{input_text}
```

---Output---
Output:
"""

messages1 = [
    SystemMessage("您是一款专用于实体提取任务的人工智能助手，请严格遵守提示指令。"),
    HumanMessage(prompt_entity_extraction1),
]

ai_msg = llm.invoke(messages1)

graph = nx.Graph()

node_map = dict()

for index, item in enumerate(ai_msg.content.split(record_delimiter)):
    line_str = item.strip()
    fields = line_str.split("|")
    if len(fields) == 4:
        print(index, fields[1])
        node_map[fields[1]] = index
        graph.add_node(index, label=fields[1], kind=fields[2], description=fields[3])

    if len(fields)==5:
        if fields[1] in node_map and fields[2] in node_map:
            print(fields)
            src_node_idx = node_map[fields[1]]
            tar_node_idx = node_map[fields[2]]
            graph.add_edge(src_node_idx, tar_node_idx, weight=1, label=fields[3], description=fields[4])

nt = Network('800px', '100%')
nt.from_nx(graph)
nt.show('gen_graph.html', notebook = False)

user_input = "刘备和张飞是什么关系？"

examples2 = f"""------Example 1------

Entity_types: [组织,人员,设备,产品,技术,地点,事件,类别]
Text:
```
while Alex clenched his jaw, the buzz of frustration dull against the backdrop of Taylor's authoritarian certainty. It was this competitive undercurrent that kept him alert, the sense that his and Jordan's shared commitment to discovery was an unspoken rebellion against Cruz's narrowing vision of control and order.

Then Taylor did something unexpected. They paused beside Jordan and, for a moment, observed the device with something akin to reverence. "If this tech can be understood..." Taylor said, their voice quieter, "It could change the game for us. For all of us."

The underlying dismissal earlier seemed to falter, replaced by a glimpse of reluctant respect for the gravity of what lay in their hands. Jordan looked up, and for a fleeting heartbeat, their eyes locked with Taylor's, a wordless clash of wills softening into an uneasy truce.

It was a small transformation, barely perceptible, but one that Alex noted with an inward nod. They had all been brought here by different paths
```

Output:
entity{tuple_delimiter}Alex{record_delimiter}
entity{tuple_delimiter}Taylor{record_delimiter}
entity{tuple_delimiter}Jordan{record_delimiter}
entity{tuple_delimiter}Cruz{record_delimiter}
entity{tuple_delimiter}The Device{record_delimiter}

"""

prompt_entity_extraction2 = f"""---Goal---
Given a text document that is potentially relevant to this activity and a list of entity types, identify all entities of those types from the text.
Use {language} as output language.

---Steps---
1. Recognizing definitively conceptualized entities in text. For each identified entity, extract the following information:
- entity_name: Name of the entity, use same language as input text. If English, capitalized the name

2. Format each entity as:
"entity"{tuple_delimiter}<entity_name>

3. Use `{tuple_delimiter}` as field delimiter. Use `{record_delimiter}` as the list delimiter. Ensure no spaces are added around the delimiters.

4. When finished, output `{completion_delimiter}`

5. Return identified entities in {language}.

---Quality Guidelines---
- Only extract entities that are clearly defined and meaningful in the context
- Avoid over-interpretation; stick to what is explicitly stated in the text
- Include specific numerical data in entity name when relevant
- Ensure entity names are consistent throughout the extraction

---Examples---
{examples2}

---Real Data---
Entity_types: [{entity_types}]
Text:
```
{user_input}
```

---Output---
Output:
"""

messages2 = [
    SystemMessage("您是一款专用于实体提取任务的人工智能助手，请严格遵守提示指令。"),
    HumanMessage(prompt_entity_extraction2),
]

query_nodes_ids = []
ai_msg2 = llm.invoke(messages2)
for index, item in enumerate(ai_msg2.content.split(record_delimiter)):
    line_str = item.strip()
    fields = line_str.split("|")
    if len(fields) == 2 and fields[1] in node_map:
        print(fields)
        query_nodes_ids.append(node_map[fields[1]])

print(query_nodes_ids)
huangjin_related = []
for u, v, attr in graph.edges(data=True):
    if u in query_nodes_ids or v in query_nodes_ids:
        relation = f"{graph.nodes[u]['label']}和{graph.nodes[v]['label']}: {attr.get('label')}, {attr.get('description')}"
        huangjin_related.append(relation)

print("huangjin_related", huangjin_related)

know_content = "\n".join(huangjin_related)

answer_prompt = f"""
基于以下已知内容回答用户问题。

回答原则:
1. 对于已知内容中包含用户问题的答案，使用简结准确的语言总结答案。
2. 当已经内容中不包含用户的问题答案时，直接回答未查询到相关内容。

已知内容: 
{know_content}

用户问题:
{user_input}

"""

messages3 = [
    SystemMessage("您是一款专用于回答问题的人工智能助手，请严格遵守提示指令。"),
    HumanMessage(answer_prompt),
]
ai_msg3 = llm.invoke(messages3)
print(ai_msg3)
